from itertools import chain
from pathlib import Path
from feasytools import TimeImplictFunc
import xml.etree.ElementTree as ET
from fpowerkit import Grid, FloatVar
from fstatistics import StaGen, StaMan, StaBus
from .base import *

def _sv(x: FloatVar)->float:
    assert x is not None, "求解失败"
    return x

class PluginPDN(PluginBase[float]):
    @property
    def Description(self)->str:
        return "配电网DistFlow模型"
    
    def Initialization(self,elem:ET.Element,inst:TrafficInst,work_dir:Path,plugin_dependency:'list[PluginBase]') -> float:
        '''从XML元素初始化插件'''
        grid_file = elem.attrib.get("grid")
        assert grid_file is not None, "未指定电网文件"
        grid_path = Path(grid_file)
        if grid_path.is_absolute(): grid_file = str(grid_path)
        else: grid_file = str(work_dir / grid_path)
        self.__gr = Grid.fromFile(grid_file,True)
        
        # 将各充电站注册至母线负荷
        for c in chain(inst.CSList,inst.PKList):
            if not c.node in self.__gr.BusNames: raise ValueError(f"错误: 充电站{c.name}的母线{c.node}在电网中不存在")
            self.__gr.Bus(c.node).Pd += TimeImplictFunc(lambda: c.Pc_kW/self.__gr.Sb_kVA)
        
        self.SetPreStep(self.PreStep)
        return 1e100

    @property
    def Grid(self)->Grid:
        '''获取电网实例'''
        return self.__gr
    
    def PreStep(self,_t:int,sta:PluginStatus)->tuple[bool,float]:
        '''求解配电网在_t时刻的最优发电计划, 求解结果请见Grid实例中的相关值'''
        if sta == PluginStatus.EXECUTE:
            return self.__gr.solve(_t)
        elif sta == PluginStatus.OFFLINE:
            return True, 0
        elif sta == PluginStatus.HOLD:
            return self.LastPostStepSucceed, self.LastPreStepResult
    
    @property
    def BestCost(self)->float: 
        '''获取最优发电计划的成本, 要求上次求解成功'''
        assert self.LastPreStepSucceed, "上次求解失败"
        return self.LastPreStepResult
    
    @property
    def GeneratorPlan(self)->dict[str,float]:
        '''获取最优发电计划, 要求上次求解成功'''
        assert self.LastPreStepSucceed, "上次求解失败"
        return {g.ID: _sv(g.P) for g in self.__gr.Gens}